package org.ojalgo.optimisation.solver.ortools;

import com.google.ortools.Loader;
import com.google.ortools.linearsolver.MPConstraint;
import com.google.ortools.linearsolver.MPObjective;
import com.google.ortools.linearsolver.MPSolver;
import com.google.ortools.linearsolver.MPVariable;
import org.ojalgo.optimisation.Expression;
import org.ojalgo.optimisation.ExpressionsBasedModel;
import org.ojalgo.optimisation.Optimisation;
import org.ojalgo.optimisation.Variable;
import org.ojalgo.structure.Structure1D.IntIndex;

import java.math.BigDecimal;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author zh
 * @date 2022/12/17 23:05
 */
public final class SolverORTools implements Optimisation.Solver {
    public static final SolverORTools.Integration INTEGRATION = new SolverORTools.Integration();
    static final SolverORTools.Configurator DEFAULT = (solver, options) -> {
    };

    static {
        Loader.loadNativeLibraries();
    }

    private final MPSolver mySolver;
    private final MPVariable[] myVariables;

    SolverORTools(MPSolver solver, MPVariable[] variables) {
        this.myVariables = variables;
        this.mySolver = solver;
    }

    static MPSolver newSolver(ExpressionsBasedModel model) {
//        return model.isAnyVariableInteger() ? MPSolver.createSolver("CBC") : MPSolver.createSolver("GLOP");
        return MPSolver.createSolver("SCIP");

    }

    static State translate(MPSolver.ResultStatus status) {
        switch (status) {
            case NOT_SOLVED:
                return State.UNEXPLORED;
//            case MODEL_INVALID:
//                return State.INVALID;
            case UNBOUNDED:
            case INFEASIBLE:
                return State.INFEASIBLE;
            case FEASIBLE:
                return State.FEASIBLE;
            case OPTIMAL:
                return State.OPTIMAL;
            case ABNORMAL:
            default:
                return State.FAILED;
        }
    }

    public void dispose() {

        if (mySolver != null) {
            mySolver.delete();
        }

    }

    public Result solve(Result kickStarter) {
        try {
            State retState = translate(mySolver.solve());
            double retValue = mySolver.objective().value();
            double[] retSolution = new double[myVariables.length];
            if (retState.isFeasible()) {
                for (int i = 0; i < retSolution.length; ++i) {
                    retSolution[i] = myVariables[i].solutionValue();
                }
            }

            return Result.of(retValue, retState, retSolution);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }

    }

    @FunctionalInterface
    public interface Configurator {
        void configure(MPSolver solver, Options options);
    }

    public static final class Integration extends ExpressionsBasedModel.Integration<SolverORTools> {
        Integration() {
        }

        public org.ojalgo.optimisation.solver.ortools.SolverORTools build(ExpressionsBasedModel model) {
            // 是否有整型变量
            boolean mip = model.isAnyVariableInteger();
            MPSolver solver = newSolver(model);

            // 参数
            List<Variable> mVars = model.getVariables();
            int nbVars = mVars.size();
            double posInf = MPSolver.infinity();
            double negInf = -posInf;
            MPVariable[] sVars = new MPVariable[nbVars];

            for (int i = 0; i < nbVars; ++i) {
                Variable mVar = mVars.get(i);
                BigDecimal lb = mVar.getLowerLimit() == null ? BigDecimal.ZERO : mVar.getLowerLimit();
                BigDecimal ub = mVar.getUpperLimit() == null ? BigDecimal.valueOf(Integer.MAX_VALUE) : mVar.getUpperLimit();
                // 整型变量
                boolean integer = mip && mVar.isInteger();
                String name = mVar.getName();
                sVars[i] = solver.makeVar(lb.doubleValue(), ub.doubleValue(), integer, name);
            }
            List<Expression> constraints = model.constraints().collect(Collectors.toList());

            model.constraints().forEach((mConstr) -> {
                BigDecimal lb = mConstr.getLowerLimit() == null ? BigDecimal.valueOf(Integer.MIN_VALUE) : mConstr.getLowerLimit();
                BigDecimal ub = mConstr.getUpperLimit() == null ? BigDecimal.valueOf(Integer.MAX_VALUE) : mConstr.getUpperLimit();
                String name = mConstr.getName();
                MPConstraint sConstr = solver.makeConstraint(lb.doubleValue(), ub.doubleValue(), name);

                for (IntIndex key : mConstr.getLinearKeySet()) {
                    sConstr.setCoefficient(sVars[key.index], mConstr.get(key).doubleValue());
                }

            });
            Expression mObj = model.objective();
            MPObjective sObj = solver.objective();
            sObj.setOptimizationDirection(model.getOptimisationSense() == Sense.MAX);

            for (IntIndex key : mObj.getLinearKeySet()) {
                sObj.setCoefficient(sVars[key.index], mObj.get(key).doubleValue());
            }

            DEFAULT.configure(solver, model.options);
            model.options.getConfigurator(org.ojalgo.optimisation.solver.ortools.SolverORTools.Configurator.class)
                    .ifPresent((c) -> c.configure(solver, model.options));
            return new org.ojalgo.optimisation.solver.ortools.SolverORTools(solver, sVars);
        }

        public boolean isCapable(ExpressionsBasedModel model) {
            return !model.isAnyExpressionQuadratic();
        }

        protected boolean isSolutionMapped() {
            return false;
        }
    }
}